// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of Binary Pedigree File (BPED)
*/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "bitstring.h"
#include "bpedfile.h"
#include "copyfile.h"
#include "locus.h"
#include "suffixes.h"
#include "types.h"
#include "utils.h"

int bpedfile_create(struct selected_options *ops, struct mapfile *m) {
	FILE *out = NULL;
	char filepath[FILE_PATH_LENGTH];

	if (ops == NULL || m == NULL) return 0;
	else if (strlen(ops->ldir) == 0 || strlen(ops->szRunId) == 0) return 0;
	else if (m->nSNP == 0 || m->nSNP == 0) return 0;

	sprintf(filepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_BPED_FILE);
	out = fopen(filepath, "wb");
	if (out == NULL) return 0;
	else {
		int i;
		register int j;
		int nSamples =  m->SNP[0]->nOne;

		fwrite(&m->nLocus, sizeof(int), 1, out); // No. SNPs
		fwrite(&nSamples, sizeof(int), 1, out); // No. samples

		for (i = 0; i < m->nLocus; i++) {
			for (j = 0; j < nSamples; j++) {
				fwrite(&m->SNP[i]->one[j], sizeof(char), 1, out);
				fwrite(&m->SNP[i]->two[j], sizeof(char), 1, out);
			}
		}
		fclose(out);
	}

	return 1;
}

int bpedfile_read_header(FILE *bin, int *nSNPs, int *nSamples) {
	int _nSamples = 0, _nSNPs = 0;
	BOOL jobDone = FALSE;

	if (bin == NULL || nSamples == NULL || nSNPs == NULL) return 0;

	if (fseek(bin, 0, SEEK_SET) == 0) {
		if (fread(&_nSNPs, sizeof(int), 1, bin) == 1) {
			if (fread(&_nSamples, sizeof(int), 1, bin) == 1) {
				if (_nSamples > 0 && _nSNPs > 0) {
					*nSNPs = _nSNPs;
					*nSamples = _nSamples;
					jobDone = TRUE;
				}
			}
		}
	}

	if (jobDone) return 1;
	else {
		*nSamples = *nSNPs = 0;
		return 0;
	}
}

char* bpedfile_fseek(FILE *bin, int nSNPs, int nSamples, int start, int end) {
	char *data = NULL;
	int header_offset = sizeof(int) * 2; // 2 integers at start of the BPED file.
	long int start_reading_from = 0, amount_to_read = 0, record_size = 0;
	BOOL jobDone = FALSE;

	if (bin == NULL || nSNPs <= 0 || nSamples <= 0 || start < 0 || end <= 0 || start > end) return NULL;

	record_size = nSamples * 2 * sizeof(char) ; // 2 bytes per allele.
	start_reading_from = (start * record_size) + header_offset;
	amount_to_read = (end - start) * record_size;
	
	data = (char*) malloc(sizeof(char) * amount_to_read);
	if (data != NULL) {
		if (fseek(bin, start_reading_from, SEEK_SET) == 0) {
			if (fread(data, amount_to_read, 1, bin) == 1)  jobDone = TRUE;
		}
	}

	if (jobDone == FALSE) {
		if (data != NULL) {
			free(data);
			data = NULL;
		}
	}

	return data;
}

int bpedfile_write_filtered(char *bpedData, int nSNPs, int nSamples, BOOL *snps_to_del, BOOL *persons_to_del, int rank, struct selected_options *ops, char *ldir, char *rdir) {
	FILE *out = NULL;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH], *buffer = NULL;
	int nActualSamples = 0, nActualSNPs = 0, i, j;
	char *pBpedData = NULL;
	BOOL jobDone = FALSE;
	size_t outer_offset = 0, inner_offset = 0, nActualRecordSize = 0;

	if (bpedData == NULL || nSNPs <= 0 || nSamples <= 0 || snps_to_del == NULL || persons_to_del == NULL || ops == NULL || rank < 0 || ldir == NULL || rdir == NULL) 
		return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ldir) == 0 || strlen(rdir) == 0) return 0;

	for (i = 0; i < nSNPs; i++) if (!snps_to_del[i]) nActualSNPs++;
	for (i = 0; i < nSamples; i++) if (!persons_to_del[i]) nActualSamples++;
	if (nActualSNPs == 0 || nActualSamples == 0) return 0;
	
	outer_offset = sizeof(char) * nSamples * 2;
	inner_offset = sizeof(char) * 2;
	nActualRecordSize = sizeof(char) * 2 * nActualSamples;

	buffer = (char*) malloc(outer_offset);
	
	sprintf(localFilepath, "%s%s_%i%s", ldir, ops->szRunId, rank, SUFFIX_FBPED_FILE);
	sprintf(remoteFilepath, "%s%s_%i%s", rdir, ops->szRunId, rank, SUFFIX_FBPED_FILE);
	out = fopen(localFilepath, "wb");
	if (buffer != NULL && out != NULL) {
		size_t amount_data = nActualSNPs * nActualSamples * 2;
		int nSNPsWrote = 0;
		
		if (fwrite(&amount_data, sizeof(size_t), 1, out) == 1) {
			for (i = 0, pBpedData = bpedData; i < nSNPs; i++, pBpedData += outer_offset) {
				char *snpRecord = pBpedData;
				char *pBuffer = buffer;

				if (snps_to_del[i]) continue;

				for (j = 0; j < nSamples; j++, snpRecord += inner_offset) {
					if (persons_to_del[j]) continue;

					pBuffer[0] = snpRecord[0];
					pBuffer[1] = snpRecord[1];
					pBuffer += inner_offset;
				}

				if (fwrite(buffer, nActualRecordSize, 1, out) == 1) nSNPsWrote++;
			}
			if (nSNPsWrote == nActualSNPs) jobDone = TRUE;
		}
	}
	if (buffer != NULL) free(buffer);
	if (out != NULL) fclose(out);
	if (jobDone) jobDone = copyAndDeleteFile(localFilepath, remoteFilepath) ? TRUE : FALSE;
	
	if (jobDone) return 1;
	else return 0;
}

size_t bpedfile_read_filtered(char *buffer, int rank, struct selected_options *ops, char *ldir, char *rdir) {
	BOOL copiedFile = FALSE;
	FILE *in = NULL;
	size_t amount_to_read = 0, amount_read = 0;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];

	if (buffer == NULL || rank < 0 || ops == NULL || ldir == NULL || rdir == NULL) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ldir) == 0 || strlen(rdir) == 0) return 0;

	// Copy the BPED fragment to local disk, then read in a single lump.
	sprintf(remoteFilepath, "%s%s_%i%s", rdir, ops->szRunId, rank, SUFFIX_FBPED_FILE);
	sprintf(localFilepath, "%s%s_%i%s", ldir, ops->szRunId, rank, SUFFIX_FBPED_FILE);
	copiedFile = copyFile(remoteFilepath, localFilepath) ? TRUE : FALSE;
	if (!copiedFile) return 0;
	
	// Open the file & read all the data into the supplied buffer.
	in = fopen(localFilepath, "rb");
	if (in != NULL) {
		if(fread(&amount_to_read, sizeof(size_t), 1, in) == 1) {
			if (amount_to_read > 0) amount_read = fread(buffer, sizeof(char), amount_to_read, in);
		}
		fclose(in);
	}
	if (copiedFile) deleteFile(localFilepath);

	return amount_read;
}

int bpedfile_write_summary_and_copy(struct selected_options *ops, struct mapfile *m) {
	BOOL jobDone = FALSE;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	int *missing_snp_counts = NULL, *missing_person_counts = NULL;

	if (ops == NULL || m == NULL) return 0;
	else if (strlen(ops->ldir) == 0 || strlen(ops->rdir) == 0 || strlen(ops->szRunId) == 0) return 0;
	else if (m->nSNP == 0 || m->SNP == NULL) return 0;
	
	missing_snp_counts = create_int_array(m->nSNP, 0);
	missing_person_counts = create_int_array(m->SNP[0]->nOne, 0);
	if (missing_snp_counts != NULL && missing_person_counts != NULL) {
		FILE *out = NULL;
		int i, nSamples =  m->SNP[0]->nOne, snp_min = -1, snp_max = -1, snp_median = -1, person_min = -1, person_max = -1, person_median = -1;
		register int j;

		for (i = 0; i < m->nLocus; i++) {
			for (j = 0; j < nSamples; j++) {
				if (m->SNP[i]->one[j] && !m->SNP[i]->two[j]) { // Allele coded as missing
					missing_snp_counts[i]++;
					missing_person_counts[j]++;
				}
			}
		}

		// Sort the counts.
		qsort(missing_snp_counts, m->nSNP, sizeof(int), int_cmp);
		qsort(missing_person_counts, nSamples, sizeof(int), int_cmp);

		// Get the median values from the sorted arrays.
		if (m->nSNP % 2 == 0) {
			i = m->nSNP / 2;
			snp_median = (missing_snp_counts[i - 1] + missing_snp_counts[i]) / 2;
		} else {
			i = (m->nSNP / 2) + 1;
			snp_median = missing_snp_counts[i];
		}
		snp_min = missing_snp_counts[0];
		snp_max = missing_snp_counts[m->nSNP - 1];

		if (nSamples % 2 == 0) {
			i = nSamples / 2;
			person_median = (missing_person_counts[i - 1] + missing_person_counts[i]) / 2;
		} else {
			i = (nSamples / 2) + 1;
			person_median = (missing_person_counts[i - 1] + missing_person_counts[i]) / 2;
		}
		person_min = missing_person_counts[0];
		person_max = missing_person_counts[nSamples - 1];
	
		// Write output to the file.
		sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_BPED_SUMMARY);
		sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_BPED_SUMMARY);
		out = fopen(localFilepath, "w");
		if (out != NULL) {
			fprintf(out, "SNP_MIN_MISSINGNESS=%1.10f\n",  (double) snp_min / (double) nSamples);
			fprintf(out, "SNP_MAX_MISSINGNESS=%1.10f\n", (double) snp_max / (double) nSamples);
			fprintf(out, "SNP_MEDIAN_MISSINGNESS=%1.10f\n", (double) snp_median / (double) nSamples);
			fprintf(out, "PERSON_MIN_MISSINGNESS=%1.10f\n", (double) person_min / (double) m->nSNP);
			fprintf(out, "PERSON_MAX_MISSINGNESS=%1.10f\n", (double) person_max / (double) m->nSNP);
			fprintf(out, "PERSON_MEDIAN_MISSINGNESS=%1.10f\n", (double) person_median / (double) m->nSNP);
			fclose(out);

			if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE;
		}
	}
	if (missing_snp_counts != NULL) free(missing_snp_counts);
	if (missing_person_counts != NULL) free(missing_person_counts);

	if (jobDone) return 1;
	else return 0;
}

// Locally temporarily appending rank to the run ID to increase the uniqueness of the output file name.
int bpedfile_create_and_copy(struct selected_options *ops, struct mapfile *m, int rank) {
	BOOL jobDone = FALSE;
	char tmp_szRunId[TOKEN_LENGTH];

	if (m == NULL || ops == NULL) return 0;

	strcpy(tmp_szRunId, ops->szRunId);
	sprintf(ops->szRunId, "%s_%i", tmp_szRunId, rank);

	if (bpedfile_create(ops, m)) {
		char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];

		sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_BPED_FILE);
		sprintf(remoteFilepath, "%s%s%s", ops->rdir, tmp_szRunId, SUFFIX_BPED_FILE);
		if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE;
	}

	strcpy(ops->szRunId, tmp_szRunId); // Restore the original Run ID

	if (jobDone) return 1;
	else return 0;
}

char* bpedfile_load(struct selected_options *ops, int *nSNPs, int *nSamples) {
	FILE *in = NULL;
	char *data = NULL;

	if (ops == NULL || nSNPs == 0 || nSamples == NULL) return data;

	in = fopen(ops->szBpedFileName, "rb");
	if (in != NULL) {
		if(bpedfile_read_header(in, nSNPs, nSamples)) {
			data = bpedfile_fseek(in, *nSNPs, *nSamples, 0, *nSNPs);
		}
		fclose(in);
	}

	return data;
}

char* _bpedfile_read_filtered(struct selected_options *ops, int rank) {
	FILE *in = NULL;
	size_t amount_to_read = 0;
	char remoteFilepath[FILE_PATH_LENGTH], *buffer = NULL;

	if (ops == NULL) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->rdir) == 0) return 0;

	sprintf(remoteFilepath, "%s%s_%i%s", ops->rdir, ops->szRunId, rank, SUFFIX_FBPED_FILE);
	in = fopen(remoteFilepath, "rb");
	if (in) {
		if(fread(&amount_to_read, sizeof(size_t), 1, in) == 1) {
			if (amount_to_read > 0){
				buffer = (char*) malloc(sizeof(char) * amount_to_read);
				if (buffer != NULL) {
					if(fread(buffer, sizeof(char), amount_to_read, in) != amount_to_read) {
						free(buffer);
						buffer = NULL;
					}
				}
			}
		}
	}
	if (in != NULL) fclose(in);
	
	return buffer;
}